package cn.johnyu;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.*;
import io.milvus.param.collection.*;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class BaseMilvusOperate {
    protected static final Logger logger = LoggerFactory.getLogger(BaseMilvusOperate.class);

    public static void main(String[] args) throws Exception {
        // 向量的维数(text-embedding-v3,可选：1024/768/512/256/128/64 的float32)
        int dim = 1024;
        String collectionName = "law_articles101";

        /* ============1. 连接到 Milvus 服务=====================*/
        logger.info("1. 正在连接到 Milvus 服务...");
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost("localhost").withPort(19530)
                .build();
        MilvusServiceClient milvusServiceClient = new MilvusServiceClient(connectParam);


        /* ==============2.检查集合是否存在=========================*/
        logger.info("2. 检查集合"+collectionName+"是否存在...");
        HasCollectionParam hasParam = HasCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();

        R<Boolean> response = milvusServiceClient.hasCollection(hasParam);
        if (response.getData()) {
            logger.info(collectionName + " 已经存在。");
        }
        // 如果不存在
        else {
            /* ===================3. 创建集合 ======================= */
            FieldType idField = FieldType.newBuilder()
                    .withName("id")
                    .withDataType(DataType.Int64)
                    .withPrimaryKey(true) // 主键(必须指定)
                    .withAutoID(true)  // 让 Milvus 自动生成主键 ID（可选）
                    .build();

            FieldType vectorField = FieldType.newBuilder()
                    .withName("embedding")
                    .withDescription("向量字段")
                    .withDataType(DataType.FloatVector)
                    .withDimension(dim)
                    .build();

            CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                    .addFieldType(idField)
                    .addFieldType(vectorField)
                    .withCollectionName(collectionName)
                    .withShardsNum(2) // 分片数
                    .build();
            //创建集合返回一个状态包装对象
            logger.info("3. 正在创建集合...");
            R<RpcStatus> rpcStatusR = milvusServiceClient.createCollection(createCollectionParam);
            logger.info("\t 创建： " + collectionName + ": " + rpcStatusR.getData().getMsg());

            //====================== 3. 创建向量的索引 ========================
            logger.info("4. 正在创建向量的索引...");
            CreateIndexParam indexParam = CreateIndexParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withFieldName("embedding")
                    .withIndexName("basic_index")
                    .withIndexType(IndexType.FLAT) //暴力全量对比
                    .withMetricType(MetricType.L2) // 计算向量之间的距离(欧氏距离)
                    .withExtraParam("{}") // 额外参数（可选）,因采用FLAT索引，所以不需要额外参数
                    .build();
            R<RpcStatus> createIndexStatusWrapper = milvusServiceClient.createIndex(indexParam);
            logger.info("\t 创建索引: " + createIndexStatusWrapper.getData().getMsg());
        }
        // =====================5. 产生随机向量列表 ============================
        logger.info("5. 正在产生随机向量列表...");
        List<List<Float>> vectors = IntStream.range(0, 3)
                .mapToObj(i -> new Random()
                        .doubles(dim)                  // 生成 double 流
                        .mapToObj(d -> (float) d)      // 转为 Float 类型
                        .collect(Collectors.toList())) // 收集成 List<Float>
                .collect(Collectors.toList());
        // ================== 6. 插入数据 =================================
        logger.info("6. 正在插入数据...");
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(Arrays.asList(
                        new InsertParam.Field("embedding", vectors)
                ))
                .build();
        R<MutationResult> insertResultWrapper = milvusServiceClient.insert(insertParam);
        logger.info("\t插入数据行数: " + insertResultWrapper.getData().getInsertCnt());

        // =================7.插入后刷新，并加载 ===========================
        logger.info("7. 插入后刷新，并加载");
        milvusServiceClient.flush(FlushParam.newBuilder()
                .withCollectionNames(Collections.singletonList(collectionName))
                .build());

        //         加载集合
        milvusServiceClient.loadCollection(LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build());

        //         等一会（或检测状态）
        Thread.sleep(1000);


        //================8. 搜索： 使用第一个向量进行搜索 ========================
        logger.info("8. 搜索： 使用第一个向量进行搜索...");
        List<Float> queryVector = vectors.get(0);

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withMetricType(MetricType.L2)
                .withOutFields(Arrays.asList("id"))
                .withTopK(3) // 返回前3个最相似的向量
                .withVectors(Collections.singletonList(queryVector)) //
                .withVectorFieldName("embedding")
                .withParams("{\"nprobe\": 10}")
                .build();
        // 执行搜索，但返回的结果比较复杂，需要借助于SearchResultsWrapper进行封装处理
        R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam);
        SearchResultData searchResultData = searchResultsR.getData().getResults();

        SearchResultsWrapper wrapper = new SearchResultsWrapper(searchResultData);

        // 假设你只搜索了一个向量（如 vectors.get(0)）
        List<SearchResultsWrapper.IDScore> results = wrapper.getIDScore(0);
        System.out.println("搜索结果如下： ");
        for (SearchResultsWrapper.IDScore result : results) {
            System.out.println("ID: " + result.getLongID() + ", Score: " + result.getScore());
        }
        // ================== 9. 关闭连接 ============================
        logger.info("9. 关闭连接...");
        milvusServiceClient.close();

    }
}
